import pymysql

# 数据库配置
db_conf = {
    "host": 'localhost',
    "user": 'root',
    "password": 'root',
    "database": 'test03',
    "charset": 'utf8mb4',
}
#约定：评论数据表（zssh_part2）id自1开始自增 ，已分析 情绪数据表（zssh_with_judge_qwen3_8b_zssh_part2）id与评论数据表id一一对应
# 表名变量定义
source_max_id_table = 'zssh_with_judge_qwen3_8b_zssh_part2'        # 用于获取最大 ID 的表-》分析好的情绪数据存的那个表
original_table = 'zssh_part2'          # 数据来源表-》评论数据表
n_parts = 1  # 分割成几个新表，默认为 2（可修改）


def get_max_id(connection, table_name):
    """获取指定表中 id 字段的最大值"""
    with connection.cursor() as cursor:
        cursor.execute(f"SELECT MAX(id) FROM `{table_name}`")
        result = cursor.fetchone()[0]
        return result if result is not None else 0


def create_partitioned_tables(connection, original_table, new_table_names):
    """根据原表结构创建多个新表"""
    with connection.cursor() as cursor:
        # 获取原表建表语句
        cursor.execute(f"SHOW CREATE TABLE `{original_table}`")
        create_sql = cursor.fetchone()[1]

        for new_table in new_table_names:
            cursor.execute(f"DROP TABLE IF EXISTS `{new_table}`")
            create_new_sql = create_sql.replace(f'`{original_table}`', f'`{new_table}`')
            cursor.execute(create_new_sql)
        connection.commit()


def split_and_insert_data(connection, original_table, new_table_names, start_id):
    """将 original_table 中 id > start_id 的数据平均插入到多个新表中"""
    with connection.cursor() as cursor:
        # 查询符合条件的数据
        cursor.execute(f"SELECT * FROM `{original_table}` WHERE id > %s ORDER BY id", (start_id,))
        rows = cursor.fetchall()

        if not rows:
            print("没有符合条件的数据需要迁移。")
            return

        total_rows = len(rows)
        print(f"共查询到 {total_rows} 条需要迁移的数据。")

        # 获取字段数量以构建插入语句
        cursor.execute(f"SELECT * FROM `{original_table}` LIMIT 1")
        col_count = len(cursor.description)
        placeholders = ','.join(['%s'] * col_count)

        # 平均分配数据到各个新表
        n_tables = len(new_table_names)
        rows_per_table = total_rows // n_tables
        remainder = total_rows % n_tables  # 多余的行数，前 remainder 个表多放一行

        current_index = 0
        for i, new_table in enumerate(new_table_names):
            # 计算当前表应插入的行数
            num_rows = rows_per_table + (1 if i < remainder else 0)
            part_data = rows[current_index:current_index + num_rows]
            current_index += num_rows

            # 构建插入语句
            insert_sql = f"INSERT INTO `{new_table}` VALUES ({placeholders})"
            cursor.executemany(insert_sql, part_data)
            print(f"已向 {new_table} 插入 {len(part_data)} 条数据。")

        connection.commit()
        print(f"✅ 数据已成功拆分并插入到 {n_parts} 个新表中。")


# 主程序
if __name__ == "__main__":
    # 生成新表名列表：zssh_part2_part1, zssh_part2_part2, ...
    new_table_names = [f"{original_table}_part{i+1}" for i in range(n_parts)]

    connection = None
    try:
        # 连接数据库
        connection = pymysql.connect(**db_conf)

        # 步骤1：获取 zssh_qwen 表的最大 ID
        max_id = get_max_id(connection, source_max_id_table)
        print(f"✅ {source_max_id_table} 表的最大 ID 为: {max_id}")

        # 步骤2：创建新表
        create_partitioned_tables(connection, original_table, new_table_names)

        # 步骤3：拆分并插入数据（id > max_id）
        split_and_insert_data(connection, original_table, new_table_names, max_id)

    except Exception as e:
        print(f"❌ 执行过程中发生错误: {e}")
        if connection:
            connection.rollback()
    finally:
        if connection:
            connection.close()